#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/point_types.h>
#include <pcl/common/transforms.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/registration/gicp.h>

#include <message_filters/subscriber.h>
#include <message_filters/synchronizer.h>
#include <message_filters/sync_policies/approximate_time.h>

#include <tf2_ros/transform_listener.h>
#include <tf2_eigen/tf2_eigen.h>
#include <Eigen/Dense>

#ifdef HAS_FAST_GICP
#include <fast_gicp/gicp/fast_gicp.hpp>
#include <fast_gicp/gicp/fast_vgicp.hpp>
#endif

using PCXYZ = pcl::PointCloud<pcl::PointXYZ>;
using PCXYZRGB = pcl::PointCloud<pcl::PointXYZRGB>;
using Policy = message_filters::sync_policies::ApproximateTime<sensor_msgs::PointCloud2, sensor_msgs::PointCloud2>;

class PcMultiMerger {
public:
  PcMultiMerger(ros::NodeHandle& nh, ros::NodeHandle& pnh) : tf_listener_(tf_buffer_) {
    pnh.param<std::string>("topic_a", topic_a_, "/rx_cloud_A");
    pnh.param<std::string>("topic_c", topic_c_, "/rx_cloud_C");
    pnh.param<std::string>("merged_topic", merged_topic_, "/merged_cloud");
    pnh.param<std::string>("fixed_frame", fixed_frame_, "map");
    pnh.param<std::string>("mode", mode_, "tf"); // "tf" | "gicp"

    pnh.param<bool>("use_vgicp", use_vgicp_, true);
    pnh.param<double>("voxel_downsample", voxel_leaf_, 0.05);
    pnh.param<int>("max_iter", max_iter_, 32);
    pnh.param<double>("max_corr_dist", max_corr_dist_, 1.0);
    pnh.param<double>("trans_eps", trans_eps_, 1e-3);
    pnh.param<double>("euclid_fitness_eps", fitness_eps_, 1e-4);
    pnh.param<double>("sync_slop", sync_slop_, 0.05);

    sub_a_.reset(new message_filters::Subscriber<sensor_msgs::PointCloud2>(nh, topic_a_, 1));
    sub_c_.reset(new message_filters::Subscriber<sensor_msgs::PointCloud2>(nh, topic_c_, 1));
    sync_.reset(new message_filters::Synchronizer<Policy>(Policy(10), *sub_a_, *sub_c_));
    sync_->setMaxIntervalDuration(ros::Duration(sync_slop_));
    sync_->registerCallback(boost::bind(&PcMultiMerger::cb, this, _1, _2));

    pub_ = nh.advertise<sensor_msgs::PointCloud2>(merged_topic_, 1);

    ROS_INFO_STREAM("pc_multi_merger mode="<<mode_
                    << " fixed_frame="<<fixed_frame_
                    << " voxel="<<voxel_leaf_
                    << " max_iter="<<max_iter_
                    << " max_corr="<<max_corr_dist_
#ifdef HAS_FAST_GICP
                    << " (fast_gicp "<<(use_vgicp_?"VGICP":"GICP")<<")"
#else
                    << " (fallback PCL::GICP)"
#endif
                    );
  }

private:
  static PCXYZ::Ptr toXYZ(const PCXYZRGB::Ptr& in){
    PCXYZ::Ptr out(new PCXYZ());
    out->reserve(in->size());
    for (auto& p : in->points) out->push_back(pcl::PointXYZ(p.x, p.y, p.z));
    return out;
  }
  static PCXYZRGB::Ptr downRGB(const PCXYZRGB::Ptr& in, double leaf){
    if (leaf <= 0.0) return in;
    pcl::VoxelGrid<pcl::PointXYZRGB> vg; vg.setLeafSize(leaf,leaf,leaf);
    vg.setInputCloud(in);
    PCXYZRGB::Ptr out(new PCXYZRGB()); vg.filter(*out);
    return out;
  }
  static PCXYZ::Ptr downXYZ(const PCXYZ::Ptr& in, double leaf){
    if (leaf <= 0.0) return in;
    pcl::VoxelGrid<pcl::PointXYZ> vg; vg.setLeafSize(leaf,leaf,leaf);
    vg.setInputCloud(in);
    PCXYZ::Ptr out(new PCXYZ()); vg.filter(*out);
    return out;
  }

  Eigen::Isometry3d tfLookup(const std::string& target, const std::string& source, const ros::Time& stamp){
    auto ts = tf_buffer_.lookupTransform(target, source, stamp, ros::Duration(0.05));
    return tf2::transformToEigen(ts.transform);
  }

  void cb(const sensor_msgs::PointCloud2ConstPtr& a_msg,
          const sensor_msgs::PointCloud2ConstPtr& c_msg) {
    PCXYZRGB::Ptr a_rgb(new PCXYZRGB()), c_rgb(new PCXYZRGB());
    pcl::fromROSMsg(*a_msg, *a_rgb);
    pcl::fromROSMsg(*c_msg, *c_rgb);

    PCXYZRGB::Ptr a_tf(new PCXYZRGB()), c_tf(new PCXYZRGB());

    if (mode_ == "tf") {
      try {
        Eigen::Isometry3d Ta = tfLookup(fixed_frame_, a_msg->header.frame_id, a_msg->header.stamp);
        Eigen::Isometry3d Tc = tfLookup(fixed_frame_, c_msg->header.frame_id, c_msg->header.stamp);
        pcl::transformPointCloud(*a_rgb, *a_tf, Ta.matrix().cast<float>());
        pcl::transformPointCloud(*c_rgb, *c_tf, Tc.matrix().cast<float>());
      } catch (const std::exception& e) {
        ROS_WARN_THROTTLE(1.0, "TF lookup failed: %s", e.what());
        return;
      }
    } else { // gicp
      PCXYZRGB::Ptr a_ds_rgb = downRGB(a_rgb, voxel_leaf_);
      PCXYZRGB::Ptr c_ds_rgb = downRGB(c_rgb, voxel_leaf_);
      PCXYZ::Ptr a_ds = toXYZ(a_ds_rgb);
      PCXYZ::Ptr c_ds = toXYZ(c_ds_rgb);

      Eigen::Matrix4f T = Eigen::Matrix4f::Identity();
      bool ok=false;

#ifdef HAS_FAST_GICP
      if (use_vgicp_) {
        fast_gicp::FastVGICP<pcl::PointXYZ,pcl::PointXYZ> vgicp;
        vgicp.setNumThreads(std::max(2u,std::thread::hardware_concurrency()));
        vgicp.setResolution(std::max(0.5f, (float)(voxel_leaf_*2.0)));
        vgicp.setMaxCorrespondenceDistance(max_corr_dist_);
        vgicp.setTransformationEpsilon(trans_eps_);
        vgicp.setEuclideanFitnessEpsilon(fitness_eps_);
        vgicp.setMaximumIterations(max_iter_);
        vgicp.setInputTarget(a_ds); vgicp.setInputSource(c_ds);
        pcl::PointCloud<pcl::PointXYZ> aligned;
        T = vgicp.align(aligned).matrix();
        ok = vgicp.hasConverged();
      } else {
        fast_gicp::FastGICP<pcl::PointXYZ,pcl::PointXYZ> gicp;
        gicp.setNumThreads(std::max(2u,std::thread::hardware_concurrency()));
        gicp.setMaxCorrespondenceDistance(max_corr_dist_);
        gicp.setTransformationEpsilon(trans_eps_);
        gicp.setEuclideanFitnessEpsilon(fitness_eps_);
        gicp.setMaximumIterations(max_iter_);
        gicp.setInputTarget(a_ds); gicp.setInputSource(c_ds);
        pcl::PointCloud<pcl::PointXYZ> aligned;
        T = gicp.align(aligned).matrix();
        ok = gicp.hasConverged();
      }
#else
      pcl::GeneralizedIterativeClosestPoint<pcl::PointXYZ,pcl::PointXYZ> gicp;
      gicp.setMaxCorrespondenceDistance(max_corr_dist_);
      gicp.setTransformationEpsilon(trans_eps_);
      gicp.setEuclideanFitnessEpsilon(fitness_eps_);
      gicp.setMaximumIterations(max_iter_);
      gicp.setInputTarget(a_ds); gicp.setInputSource(c_ds);
      pcl::PointCloud<pcl::PointXYZ> aligned;
      gicp.align(aligned);
      ok = gicp.hasConverged();
      T = gicp.getFinalTransformation();
#endif
      if (!ok) { ROS_WARN_THROTTLE(1.0, "GICP not converged"); return; }

      pcl::transformPointCloud(*a_rgb, *a_tf, Eigen::Matrix4f::Identity());
      pcl::transformPointCloud(*c_rgb, *c_tf, T);
    }

    PCXYZRGB::Ptr merged(new PCXYZRGB());
    merged->reserve(a_tf->size()+c_tf->size());
    *merged += *a_tf; *merged += *c_tf;

    sensor_msgs::PointCloud2 out;
    pcl::toROSMsg(*merged, out);
    out.header.frame_id = (mode_=="tf" ? fixed_frame_ : a_msg->header.frame_id);
    out.header.stamp = std::max(a_msg->header.stamp, c_msg->header.stamp);
    pub_.publish(out);
  }

  // 参数
  std::string topic_a_, topic_c_, merged_topic_, fixed_frame_, mode_;
  bool use_vgicp_{true};
  double voxel_leaf_{0.05}, max_corr_dist_{1.0}, trans_eps_{1e-3}, fitness_eps_{1e-4}, sync_slop_{0.05};
  int max_iter_{32};

  // ROS
  std::unique_ptr<message_filters::Subscriber<sensor_msgs::PointCloud2>> sub_a_, sub_c_;
  std::unique_ptr<message_filters::Synchronizer<Policy>> sync_;
  ros::Publisher pub_;
  tf2_ros::Buffer tf_buffer_;
  tf2_ros::TransformListener tf_listener_;
};

int main(int argc, char** argv) {
  ros::init(argc, argv, "pc_multi_merger");
  ros::NodeHandle nh, pnh("~");
  PcMultiMerger node(nh, pnh);
  ros::spin();
  return 0;
}

